package com.andova.camerax

import android.media.Image
import java.nio.ByteBuffer
import kotlin.math.min

class ImagePlane {
    private var mNV21: ByteArray? = null

    @Suppress("FunctionName")
    fun YUVToNV21(image: Image?): ByteArray? {
        image ?: return null
        val planes = image.planes
        if (planes.size != 3) return null
        val yPlane = planes[0]
        val uPlane = planes[1]
        val vPlane = planes[2]
        val yBuffer: ByteBuffer = yPlane.buffer
        val uBuffer: ByteBuffer = uPlane.buffer
        val vBuffer: ByteBuffer = vPlane.buffer
        yBuffer.rewind()
        uBuffer.rewind()
        vBuffer.rewind()
        val ySize = yBuffer.remaining()
        var position = 0

        val width = image.width
        val height = image.height
        if (mNV21 == null) mNV21 = ByteArray(ySize + width * height / 2)

        // Add the full y buffer to the array. If rowStride > 1, some padding may be skipped.
        for (row in 0 until height) {
            yBuffer[mNV21!!, position, width]
            position += width
            yBuffer.position(min(ySize, yBuffer.position() - width + yPlane.rowStride))
        }
        val chromaHeight: Int = height / 2
        val chromaWidth: Int = width / 2
        val vRowStride: Int = vPlane.rowStride
        val uRowStride: Int = uPlane.rowStride
        val vPixelStride: Int = vPlane.pixelStride
        val uPixelStride: Int = uPlane.pixelStride

        // Interleave the u and v frames, filling up the rest of the buffer. Use two line buffers to
        // perform faster bulk gets from the byte buffers.
        val vLineBuffer = ByteArray(vRowStride)
        val uLineBuffer = ByteArray(uRowStride)
        for (row in 0 until chromaHeight) {
            vBuffer[vLineBuffer, 0, min(vRowStride, vBuffer.remaining())]
            uBuffer[uLineBuffer, 0, min(uRowStride, uBuffer.remaining())]
            var vLineBufferPosition = 0
            var uLineBufferPosition = 0
            for (col in 0 until chromaWidth) {
                mNV21!![position++] = vLineBuffer[vLineBufferPosition]
                mNV21!![position++] = uLineBuffer[uLineBufferPosition]
                vLineBufferPosition += vPixelStride
                uLineBufferPosition += uPixelStride
            }
        }
        return mNV21
    }
}